import numpy as np
import cv2
import math
from scipy import ndimage

img = cv2.imread("6.jpg")
img1 = img.copy()
cv2.imshow("img", img)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# 使用Canny边缘检测算法进行边缘检测
edges = cv2.Canny(gray, 100, 150, apertureSize=3)
lines = cv2.HoughLines(edges, 1, np.pi / 180, 110)
# 初始化旋转角度为0
rotate_angle = 0
# 遍历检测到的直线
for rho, theta in lines[0]:
    a = np.cos(theta)
    b = np.sin(theta)
    x0 = a * rho
    y0 = b * rho
    x1 = int(x0 + 1000 * (-b))
    y1 = int(y0 + 1000 * (a))
    x2 = int(x0 - 1000 * (-b))
    y2 = int(y0 - 1000 * (a))
    # 排除垂直或水平的直线
    if x1 == x2 or y1 == y2:
        continue
    # 计算直线的斜率
    t = float(y2 - y1) / (x2 - x1)
    # 根据斜率计算旋转角度
    rotate_angle = math.degrees(math.atan(t))
    print(rotate_angle)
    rotate_angle = rotate_angle*2
    # 根据旋转角度进行调整
    if rotate_angle > 45:
        rotate_angle = -90 + rotate_angle
    elif rotate_angle < -45:
        rotate_angle = 90 + rotate_angle

for line in lines:
    rho = line[0][0]
    theta = line[0][1]
    a = np.cos(theta)
    b = np.sin(theta)
    x0 = a * rho
    y0 = b * rho
    x1 = int(x0 + 1000 * (-b))
    y1 = int(y0 + 1000 * (a))
    x2 = int(x0 - 1000 * (-b))
    y2 = int(y0 - 1000 * (a))
    cv2.line(img1, (x1, y1), (x2, y2), (0, 0, 255), 1)
cv2.imshow("lines", img1)
rotate_img = ndimage.rotate(img, rotate_angle)
cv2.imshow("ro-14-rotate_img", rotate_img)
cv2.waitKey(0)
